import os
import shutil
import random
from PIL import Image, ImageEnhance

def count_photos_in_subfolders(parent_folder):
    subfolders = [f.path for f in os.scandir(parent_folder) if f.is_dir()]
    photo_counts = {}

    for subfolder in subfolders:
        photos = [f for f in os.listdir(subfolder) if f.endswith(('.png', '.jpg', '.jpeg'))]
        photo_counts[subfolder] = len(photos)

    return photo_counts

def enhance_photo(photo_path, i):
    image = Image.open(photo_path)
    enhancer = ImageEnhance.Brightness(image)
    enhanced_image = enhancer.enhance(1.5)  # Increase brightness
    enhanced_image = enhanced_image.rotate(30)  # Rotate image
    enhanced_image.save(os.path.join(os.path.dirname(photo_path), f'copy_{os.path.basename(photo_path).replace(".", f"_enhanced_{i}.")}'))

def move_photos_to_new_folder(parent_folder, new_folder, photo_counts):
    total_photos = sum(photo_counts.values())
    avg_photo_count = total_photos // len(photo_counts)
    print(f'Total photos: {total_photos}, Average photos: {avg_photo_count}')

    for subfolder, count in photo_counts.items():
        photos = [f for f in os.listdir(subfolder) if f.endswith(('.png', '.jpg', '.jpeg'))]
        print(f'Subfolder: {subfolder}, Photos: {count}')

        new_subfolder_path = os.path.join(new_folder, os.path.basename(subfolder))
        if not os.path.exists(new_subfolder_path):
            os.makedirs(new_subfolder_path)

        if count > avg_photo_count:
            selected_photos = random.sample(photos, avg_photo_count)
            for photo in selected_photos:
                shutil.move(os.path.join(subfolder, photo), new_subfolder_path)
        else:
            for photo in photos:
                shutil.move(os.path.join(subfolder, photo), new_subfolder_path)
            if photos:  # Check if there are photos in the subfolder
                for i in range(avg_photo_count - count):
                    photo = random.choice(photos)
                    enhance_photo(os.path.join(new_subfolder_path, photo), i)
                    print(f'Enhanced photos: {i+1}')

parent_folder = 'train_more'
new_folder = 'train'

photo_counts = count_photos_in_subfolders(parent_folder)
move_photos_to_new_folder(parent_folder, new_folder, photo_counts)
